import numpy as np
from os import listdir
import functools
import math


# 文件里读出内容转换为numpy向量
def load_file_to_vector(filename):
    returnVector = np.zeros((1, 1024))

    # 打开文件
    f = open(filename)

    # 文件均为32 * 32的01矩阵
    for i in range(32):
        lineStr = f.readline()

        for j in range(32):
            returnVector[0, 32 * i + j] = int(lineStr[j])

    return returnVector


# load trainData
def load_train_dataset():
    # 训练集存放的文件名称
    dictionary_name = "knn_work/trainingDigits"
    train_file_list = listdir(dictionary_name)
    # 拿到文件目录下的文件数量
    file_count = len(train_file_list)

    # 训练集矩阵
    train_dataset = np.zeros((file_count, 1024))
    # 训练集的labels
    train_labels = []

    # 开始load数据
    print("loading train_dataset now .....")
    for i in range(file_count):
        # 文件名称 label_count.txt
        file_name = train_file_list[i]
        # 获得label
        label = int(file_name.split('_')[0])
        train_labels.append(label)
        # 将文件变成矩阵
        train_dataset[i, :] = load_file_to_vector(dictionary_name + '/' + file_name)

    # 测试
    print("load finished, size : " + str(len(train_labels)))
    return train_labels, train_dataset


# 计算两个向量的欧氏距离
def count_dis(vec1, vec2):
    sum = 0
    for i in range(1024):
        if int(vec1[0][i]) - int(vec2[i]) != 0:
            sum += 1
    return math.sqrt(sum)


# cmp, 用于get_k_min_dis_label的list排序
def cmp(obj1, obj2):
    dis1, dis2 = obj1['dis'], obj2['dis']
    if dis1 > dis2:
        return 1
    elif dis1 < dis2:
        return -1
    else:
        return 0


# 计算离目标矩阵最近的k个label
def get_k_min_dis_label(input_vector, train_dict, k):
    # 存储结果的list，其大小固定为k
    k_label = []
    # 存储所有距离的list, 里面存放的是字典对象 {'index', 'dis'}
    dis_list = []
    # 开始计算目标向量对所有训练集的距离
    for i in range(len(train_dict['dataset'])):
        dis_dict = {'index': i, 'dis': count_dis(input_vector, train_dict['dataset'][i])}
        dis_list.append(dis_dict)

    # 对dict_list排序，排序依据为字典对象内的dis
    dis_list.sort(key=functools.cmp_to_key(cmp))
    # 最小距离的k个目标
    k_res = []
    for i in range(k):
        label = train_dict['labels'][dis_list[i]['index']]
        k_res.append(label)
    return k_res


# 从k个目标label中得出最终结论
def from_k_get_label(k_label):
    # k个label中的实体数量
    label_count = {'0': 0, '1': 0, '2': 0, '3': 0, '4': 0, '5': 0, '6': 0, '7': 0, '8': 0, '9': 0}

    # 取出标记
    for i in range(len(k_label)):
        index = str(k_label[i])
        label_count[index] = label_count[index] + 1

    # 统计最多的label
    max_count_label = 0
    for i in range(10):
        if label_count[str(i)] > label_count[str(max_count_label)]:
            max_count_label = i
    print(label_count)
    return max_count_label


# knn训练
def knn(input, train_dict, k=3):
    return from_k_get_label(get_k_min_dis_label(input, train_dict, k))


# 测试集验证
def test(train_dict):
    # 训练集存放的文件名称
    dictionary_name = "knn_work/testDigits"
    test_file_list = listdir(dictionary_name)
    # 拿到文件目录下的文件数量
    file_count = len(test_file_list)

    # 正确的数量
    num = 0

    # 开始验证
    print("testing test_dataset now .....")
    for i in range(file_count):
        # 文件名称 label_count.txt
        file_name = test_file_list[i]
        # 获得label
        label = int(file_name.split('_')[0])
        # 将文件变成矩阵
        test_data = load_file_to_vector(dictionary_name + '/' + file_name)
        train_label = knn(test_data, train_dict, 5)
        if label == train_label:
            num += 1
        else:
            print("识别错误，样本 ： ", file_name, " 识别结果: ", train_label)
    print("正确率 ： ", num / file_count)


def work(num, vectors) :
    # 获得数据集
    labels, dataset = load_train_dataset()

    # 训练集数据的字典
    train_dict = {'labels': labels, 'dataset': dataset}

    # knn训练参数
    k = 5

    result = []

    for i in range(num):
        ans = str(knn(vectors[i], train_dict, k))
        print("第" + str(i) + "个数字是 ： " + str(ans))
        result.append(ans)

    return result


if __name__ == '__main__':
    # 获得数据集
    labels, dataset = load_train_dataset()

    # 训练集数据的字典
    train_dict = {'labels': labels, 'dataset': dataset}

    # # 验证集验证
    # test(train_dict)

    # knn训练参数
    k = 5
    input_file_name = "trainingDigits/5_101.txt"
    input_vector = load_file_to_vector(input_file_name)

    # 输出结果
    print("the answer is : " + str(knn(input_vector, train_dict, k)))